package com.example.springaiollamademo.websocket;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import com.example.springaiollamademo.service.ChatService;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;

/**
 * WebSocket聊天处理器
 * @author 张飞
 * @version 1.0
 * @since 2025-12-02
 */
@Component
public class ChatWebSocketHandler extends TextWebSocketHandler {

    private static final Logger logger = LoggerFactory.getLogger(ChatWebSocketHandler.class);
    
    private final ChatService chatService;
    private final ObjectMapper objectMapper;
    
    // 存储所有活动的WebSocket会话
    private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();

    public ChatWebSocketHandler(ChatService chatService) {
        this.chatService = chatService;
        this.objectMapper = new ObjectMapper();
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        logger.info("=== WebSocket连接建立 ===");
        logger.info("会话ID: {}", session.getId());
        sessions.put(session.getId(), session);
        
        // 发送连接成功消息
        sendMessage(session, "connected", "WebSocket连接成功");
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        logger.info("=== 收到WebSocket消息 ===");
        logger.info("会话ID: {}, 消息: {}", session.getId(), message.getPayload());
        
        try {
            // 解析客户端消息
            JsonNode jsonNode = objectMapper.readTree(message.getPayload());
            String prompt = jsonNode.get("prompt").asText();
            String systemMessage = jsonNode.has("systemMessage") ? 
                jsonNode.get("systemMessage").asText() : "";
            
            logger.info("解析参数: prompt=[{}], systemMessage=[{}]", prompt, systemMessage);
            
            // 发送"思考中"提示
            sendMessage(session, "thinking", "AI正在思考中，请稍候...");
            
            // 调用流式聊天服务
            chatService.chatWithSystemStream(prompt, systemMessage, new ChatService.StreamCallback() {
                @Override
                public void onChunk(String chunk) {
                    try {
                        // 发送数据块
                        sendMessage(session, "chunk", chunk);
                    } catch (Exception e) {
                        logger.error("发送数据块失败", e);
                    }
                }

                @Override
                public void onComplete() {
                    try {
                        // 发送完成消息
                        sendMessage(session, "done", "流式响应完成");
                    } catch (Exception e) {
                        logger.error("发送完成消息失败", e);
                    }
                }

                @Override
                public void onError(Throwable error) {
                    try {
                        // 发送错误消息
                        sendMessage(session, "error", "处理失败: " + error.getMessage());
                    } catch (Exception e) {
                        logger.error("发送错误消息失败", e);
                    }
                }
            });
            
        } catch (Exception e) {
            logger.error("处理WebSocket消息失败", e);
            sendMessage(session, "error", "处理消息失败: " + e.getMessage());
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        logger.info("=== WebSocket连接关闭 ===");
        logger.info("会话ID: {}, 状态: {}", session.getId(), status);
        sessions.remove(session.getId());
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        logger.error("WebSocket传输错误", exception);
        sessions.remove(session.getId());
    }

    /**
     * 发送消息到客户端
     */
    private void sendMessage(WebSocketSession session, String type, String data) throws IOException {
        if (session.isOpen()) {
            String jsonMessage = String.format("{\"type\":\"%s\",\"data\":\"%s\"}", 
                type, escapeJson(data));
            session.sendMessage(new TextMessage(jsonMessage));
            logger.debug("发送消息: type={}, data={}", type, data);
        }
    }

    /**
     * 转义JSON字符串
     */
    private String escapeJson(String str) {
        if (str == null) return "";
        return str.replace("\\", "\\\\")
                  .replace("\"", "\\\"")
                  .replace("\n", "\\n")
                  .replace("\r", "\\r")
                  .replace("\t", "\\t");
    }
}
